// ============================================================================
// File: cmd/onyx-snapshotter/main.go
// ============================================================================
package main

import (
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"

	snapshotsapi "github.com/containerd/containerd/api/services/snapshots/v1"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/config"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/metrics"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/recovery"
	"gitee.com/wangtsingx/onyx-snapshotter/pkg/snapshotter"
)

var (
	configPath = flag.String("config", "/etc/onyx/config.yaml", "path to config file")
	address    = flag.String("address", "/run/onyx/onyx.sock", "snapshotter socket address")
	logLevel   = flag.String("log-level", "info", "log level")
	metricsAddr = flag.String("metrics-addr", ":9090", "metrics endpoint address")
)

func main() {
	flag.Parse()

	// Setup logging
	level, err := logrus.ParseLevel(*logLevel)
	if err != nil {
		logrus.WithError(err).Fatal("failed to parse log level")
	}
	logrus.SetLevel(level)
	logrus.SetFormatter(&logrus.JSONFormatter{})

	// Load configuration
	cfg, err := config.Load(*configPath)
	if err != nil {
		logrus.WithError(err).Fatal("failed to load config")
	}

	// Initialize metrics
	metrics.Init()

	// Try to recover from previous crash
	if err := recovery.RecoverSnapshots(cfg.Root); err != nil {
		logrus.WithError(err).Warn("failed to recover snapshots")
	}

	// Initialize snapshotter
	snap, err := snapshotter.NewSnapshotter(cfg)
	if err != nil {
		logrus.WithError(err).Fatal("failed to create snapshotter")
	}
	defer snap.Close()

	// Start metrics server
	go func() {
		http.Handle("/metrics", promhttp.Handler())
		if err := http.ListenAndServe(*metricsAddr, nil); err != nil {
			logrus.WithError(err).Error("metrics server failed")
		}
	}()

	// Remove existing socket
	if err := os.RemoveAll(*address); err != nil {
		logrus.WithError(err).Fatal("failed to remove existing socket")
	}

	// Create gRPC server with interceptors
	rpc := grpc.NewServer(
		grpc.UnaryInterceptor(metrics.UnaryServerInterceptor()),
	)
	service := snapshotter.NewService(snap)
	snapshotsapi.RegisterSnapshotsServer(rpc, service)

	// Listen on socket
	listener, err := net.Listen("unix", *address)
	if err != nil {
		logrus.WithError(err).Fatal("failed to listen")
	}
	defer listener.Close()

	// Handle signals
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		sig := <-sigCh
		logrus.Infof("received signal %v, shutting down...", sig)
		rpc.GracefulStop()
	}()

	logrus.Infof("onyx-snapshotter v2.0 listening on %s", *address)
	logrus.Infof("metrics available at http://localhost%s/metrics", *metricsAddr)
	
	if err := rpc.Serve(listener); err != nil {
		logrus.WithError(err).Fatal("failed to serve")
	}
}